# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""数据流分析 Pass

为每个函数构建数据流信息：
- Def-Use Chain: 变量的定义和使用关系
- 逃逸分析: 变量是否逃逸到外部作用域
- 生命周期分析: 变量的生命周期范围
- 异步函数分析: Segment 划分、最小代码单元划分、变量提升

这些信息用于：
- Lambda 存储模式决策（栈分配 vs 堆分配）
- 闭包变量捕获策略
- 异步函数代码生成（变量提升、co_yield 生成）
- 未来的编译器优化（常量传播、死代码消除等）
"""

from typing import Dict, Optional, Any
from dataclasses import dataclass, field
from mcpy.analysis.dataflow import (
    DataFlowInfo,
    Definition,
    Use,
    UseType,
    EscapeInfo,
    ClosureVariableUsage,
)
from mcpy.passes.base import BasePass
from mcpy.ir.expressions.lambda_expr import IRLambda
from mcpy.ir.statements.assignment import IRAssign
from mcpy.ir.expressions.basic import IRName
from mcpy.ir.declarations.function import IRFunctionNode


class DataFlowAnalysisPass(BasePass):
    """数据流分析 Pass

    为每个函数构建数据流信息，存储到函数节点的 'dataflow' 扩展参数中
    """

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__()

    def run(self, ir_module, context):
        """执行数据流分析

        Args:
            ir_module: IR 模块
            context: Pass 上下文

        Returns:
            IR 模块（未修改）
        """
        # 遍历所有函数
        for node in ir_module.walk():
            if isinstance(node, IRFunctionNode):
                self._analyze_function(node)

        return ir_module

    def _analyze_function(self, func) -> None:
        """分析单个函数的数据流

        Args:
            func: 函数节点
        """
        # 检查函数是否有 body
        if not hasattr(func, "body") or func.body is None:
            return

        if not hasattr(func.body, "statements"):
            return

        # 创建数据流信息
        dataflow = DataFlowInfo()

        # 步骤1: 收集定义和使用
        for i, stmt in enumerate(func.body.statements):
            self._collect_defs_and_uses(stmt, i, dataflow)

        # 步骤2: 执行逃逸分析
        self._analyze_escapes(dataflow)

        # 步骤3: 闭包变量使用分析
        self._analyze_closure_variables(func, dataflow)

        # 步骤4: Lambda 存储模式决策（基于逃逸分析）
        self._determine_lambda_storage_modes(func, dataflow)

        # 步骤5: 异步函数的特殊数据流分析
        if func.is_async:
            has_await = func.get_extension_param("has_await_expressions", False)
            if has_await:
                self._analyze_async_dataflow(func, dataflow)

        # 步骤6: 存储到函数节点
        func.set_extension_param("dataflow", dataflow)

    def _collect_defs_and_uses(
        self, stmt, stmt_index: int, dataflow: DataFlowInfo
    ) -> None:
        """收集语句中的变量定义和使用

        Args:
            stmt: IR 语句节点
            stmt_index: 语句在函数中的索引
            dataflow: 数据流信息
        """
        # 遍历语句中的所有节点
        for node in stmt.walk():
            # 收集定义点：赋值语句的目标
            if isinstance(node, IRAssign):
                target = node.get_target_expr()
                if isinstance(target, IRName):
                    definition = Definition(
                        var_name=target.id,
                        node=node,
                        statement=stmt,
                        stmt_index=stmt_index,
                    )
                    dataflow.add_definition(target.id, definition)

            # 收集使用点：变量名（不在赋值左侧）
            elif isinstance(node, IRName):
                if not self._is_assignment_target(node):
                    # 检查是否是类型相关的名称（不应该被识别为变量）
                    if self._is_type_name(node):
                        continue  # 跳过类型名称
                    use = Use(
                        var_name=node.id,
                        node=node,
                        statement=stmt,
                        stmt_index=stmt_index,
                        use_type=self._classify_use_type(node),
                    )
                    dataflow.add_use(node.id, use)

    def _is_assignment_target(self, name_node) -> bool:
        """判断变量名是否是赋值目标（定义点）

        Args:
            name_node: 变量名节点

        Returns:
            True 表示是赋值目标
        """
        from mcpy.ir.statements.assignment import IRAssign

        parent = name_node.parent
        if isinstance(parent, IRAssign):
            target = parent.get_target_expr()
            return target == name_node
        return False

    def _is_type_name(self, name_node: IRName) -> bool:
        """判断 IRName 是否是类型名称（不应该被识别为变量）

        以下情况应该被识别为类型名称：
        1. IRName 是 IRSubscript 的一部分（如 list[int] 中的 list 和 int），
           且该 IRSubscript 是 IRCall 的 func（类型构造函数调用）
        2. IRName 是 IRCall 的 func，且该 IRCall 是类型构造函数调用（如 int()）
        3. IRName 在类型注解中

        Args:
            name_node: IRName 节点

        Returns:
            True 表示是类型名称，不应该被识别为变量
        """
        from mcpy.ir.expressions.access import IRSubscript
        from mcpy.ir.expressions.call import IRCall
        from mcpy.symbol_meta import SymbolKind

        parent = name_node.parent
        if parent is None:
            return False

        # 情况1: IRName 是 IRSubscript 的一部分（如 list[int] 中的 list 和 int）
        if isinstance(parent, IRSubscript):
            # 检查 IRSubscript 是否是 IRCall 的 func（类型构造函数调用）
            subscript_parent = parent.parent
            if isinstance(subscript_parent, IRCall):
                # list[int]() 这种类型构造函数调用
                if subscript_parent.func == parent:
                    return True

            # 检查 IRSubscript 是否在类型注解中
            # 向上遍历查找，看是否在类型注解上下文中
            current = subscript_parent if subscript_parent else parent
            depth = 0
            while current and depth < 5:  # 限制遍历深度
                # 如果在类型注解中，也应该跳过
                # 这里可以根据需要扩展
                if hasattr(current, "node_type") and current.node_type in [
                    "annotation",
                    "type_alias",
                ]:
                    return True
                current = current.parent if hasattr(current, "parent") else None
                depth += 1

        # 情况2: IRName 本身是 IRCall 的 func（如 int()）
        if isinstance(parent, IRCall) and parent.func == name_node:
            # 检查是否是类型构造函数调用
            # 通过符号查找判断是否是类型符号
            try:
                symbol = name_node.lookup_symbol(
                    name_node.id, symbol_kinds=[SymbolKind.CLASS]
                )
                if symbol and symbol.kind == SymbolKind.CLASS:
                    # 是类型符号，且是函数调用的 func，说明是类型构造函数调用
                    return True
            except Exception:
                # 符号查找可能失败，忽略
                pass

        return False

    def _classify_use_type(self, name_node) -> UseType:
        """分类变量使用类型

        Args:
            name_node: 变量名节点

        Returns:
            使用类型
        """
        # 动态导入
        from mcpy.ir.expressions.call import IRCall
        from mcpy.ir.statements.assignment import IRAssign
        from mcpy.ir.statements.control import IRReturn
        from mcpy.ir.expressions.literals import IRList, IRDict, IRSet

        parent = name_node.parent
        if not parent:
            return UseType.OTHER

        # 场景1: 函数调用 f(args)
        if isinstance(parent, IRCall) and parent.func == name_node:
            return UseType.CALL

        # 场景2: 二次赋值 g = f
        if isinstance(parent, IRAssign) and parent.value == name_node:
            return UseType.ASSIGNMENT

        # 场景3: 返回 return f
        if isinstance(parent, IRReturn):
            return UseType.RETURN

        # 场景4: 容器字面量 [f] 或 {f}
        if isinstance(parent, (IRList, IRDict, IRSet)):
            return UseType.CONTAINER

        # 场景5: 作为参数传递 func(f)
        if isinstance(parent, IRCall):
            # parent.func != name_node 已在场景1排除
            return UseType.PARAMETER

        return UseType.OTHER

    def _analyze_escapes(self, dataflow: DataFlowInfo) -> None:
        """分析变量逃逸

        逃逸规则：
        - 二次赋值 (g = f) → 逃逸（需要拷贝/共享所有权）
        - 返回 (return f) → 逃逸（生命周期延长）
        - 存储到容器 ([f]) → 逃逸（持久化存储）
        - 作为参数传递 → 暂时认为逃逸（保守策略）

        Args:
            dataflow: 数据流信息
        """
        for var_name, uses in dataflow.uses.items():
            escape_info = EscapeInfo()

            for use in uses:
                # 逃逸规则检查
                if use.use_type in (
                    UseType.ASSIGNMENT,
                    UseType.RETURN,
                    UseType.CONTAINER,
                    UseType.PARAMETER,  # 保守：参数传递暂时认为逃逸
                ):
                    escape_info.escapes = True
                    escape_info.reason = (
                        f"{use.use_type.value} at stmt {use.stmt_index}"
                    )
                    escape_info.escape_sites.append(use)
                    break  # 一旦发现逃逸，无需继续检查

            dataflow.escape_info[var_name] = escape_info

    def _analyze_closure_variables(self, func, dataflow: DataFlowInfo) -> None:
        """分析闭包变量的使用模式（支持 Lambda 和嵌套函数）

        检测：
        1. 变量是否在 lambda/嵌套函数 定义后被修改
        2. 变量是否被多个 lambda/嵌套函数 共享

        Args:
            func: 函数节点
            dataflow: 数据流信息
        """

        # 收集函数中所有的 lambda 和嵌套函数
        closures = []  # 可以是 IRLambda 或 IRFunctionNode
        for node in func.walk():
            if isinstance(node, IRLambda):
                closures.append(node)
            elif isinstance(node, IRFunctionNode) and node != func:
                # 嵌套函数（不是当前函数本身）
                if self._is_nested_function_of(node, func):
                    closures.append(node)

        if not closures:
            return  # 没有 lambda 或嵌套函数，无需分析

        # 为每个闭包（lambda 或嵌套函数）分析其捕获的变量
        for closure_node in closures:
            # 获取捕获变量列表（lambda 和嵌套函数都使用相同的属性）
            # Lambda 和嵌套函数都使用 captured_vars 属性
            captured_vars = closure_node.captured_vars or []

            # 如果闭包是嵌套函数且 captured_vars 为空，补充设置简单函数标记
            if isinstance(closure_node, IRFunctionNode):
                if not captured_vars:
                    # 检查标记是否已设置（可能在 TypeInferenceEngine 中已设置）
                    is_simple = closure_node.get_extension_param(
                        "is_simple_non_capturing", False
                    )
                    if not is_simple:
                        # 补充设置标记
                        closure_node.set_extension_param("is_simple_non_capturing", True)

            for var_name in captured_vars:
                # 为每个闭包变量创建使用信息
                usage = self._analyze_closure_var_usage(
                    var_name, closure_node, func, closures, dataflow
                )

                # 存储到节点（每个闭包独立存储）
                closure_node.set_extension_param(f"closure_usage_{var_name}", usage)

    def _is_nested_function_of(
        self, nested_func: IRFunctionNode, outer_func: IRFunctionNode
    ) -> bool:
        """判断 nested_func 是否是 outer_func 的直接嵌套函数

        Args:
            nested_func: 可能的嵌套函数
            outer_func: 外层函数

        Returns:
            bool: 是否是直接嵌套函数
        """
        parent = nested_func.parent
        while parent:
            if parent == outer_func:
                return True
            if isinstance(parent, IRFunctionNode):
                # 遇到其他函数，说明不是直接嵌套
                return False
            parent = parent.parent
        return False

    def _analyze_closure_var_usage(
        self,
        var_name: str,
        closure_node,
        func,
        all_closures: list,
        dataflow: DataFlowInfo,
    ) -> ClosureVariableUsage:
        """分析单个闭包变量的使用（支持 Lambda 和嵌套函数）

        Args:
            var_name: 闭包变量名
            closure_node: Lambda 或嵌套函数节点
            func: 包含函数
            all_closures: 函数中所有闭包（Lambda 和嵌套函数）列表
            dataflow: 数据流信息

        Returns:
            闭包变量使用信息
        """
        usage = ClosureVariableUsage()

        # 检查变量是否是外层函数的参数
        var_symbol = func.lookup_symbol(var_name)
        if var_symbol:
            from mcpy.symbol_meta import SymbolKind

            usage.is_parameter = var_symbol.kind == SymbolKind.PARAMETER

        # 获取闭包定义所在的语句索引
        closure_stmt_index = self._get_closure_stmt_index(closure_node, func)
        if closure_stmt_index is None:
            return usage  # 无法确定位置，返回默认值

        # 问题1：变量是否在闭包定义后被修改？
        # 检查定义点（赋值目标），而非使用点
        var_defs = dataflow.definitions.get(var_name, [])
        for defn in var_defs:
            # 只检查闭包定义之后的定义
            if defn.stmt_index > closure_stmt_index:
                # 闭包之后有新的定义 → 变量被修改
                usage.modified_after_capture = True
                break  # 发现修改，无需继续

        # 问题2：变量是否被多个闭包捕获？
        capture_count = 0
        for closure in all_closures:
            # 获取捕获变量列表
            # Lambda 和嵌套函数都使用 captured_vars 属性
            captured = closure.captured_vars or []

            if var_name in captured:
                capture_count += 1

        if capture_count > 1:
            usage.shared_by_multiple_lambdas = True

        # 问题3：生命周期是否超过闭包？（预留）
        # 简化：如果闭包逃逸，则闭包变量生命周期超过闭包
        # 未来可以更精确分析

        return usage

    def _get_closure_stmt_index(self, closure_node, func) -> int | None:
        """获取闭包（Lambda 或嵌套函数）定义所在的语句索引

        Args:
            closure_node: Lambda 或函数节点
            func: 包含函数

        Returns:
            语句索引，如果找不到返回 None
        """
        # 向上追溯到语句节点
        current = closure_node
        while current and current.parent != func.body:
            current = current.parent

        if not current:
            return None

        # 在 func.body.statements 中查找索引
        try:
            return func.body.statements.index(current)
        except (ValueError, AttributeError):
            return None

    def _determine_lambda_storage_modes(self, func, dataflow: DataFlowInfo) -> None:
        """为函数中的所有 Lambda 决定存储模式（栈 vs 堆）

        决策规则：
        1. 简单赋值场景（赋值给局部变量且不逃逸）→ 栈模式
        2. 其他场景（传参、返回、容器、逃逸）→ 堆模式

        注意：栈模式也支持默认参数（通过 set_defaults）

        Args:
            func: 函数节点
            dataflow: 数据流信息
        """

        # 遍历函数中的所有 lambda
        for node in func.walk():
            if not isinstance(node, IRLambda):
                continue

            # 默认：堆模式
            storage_mode = "heap"

            # 检查是否是简单赋值场景
            # 条件：lambda 直接赋值给局部变量，且该变量不逃逸
            parent = node.parent
            if isinstance(parent, IRAssign):
                target = parent.get_target_expr()
                if isinstance(target, IRName):
                    var_name = target.id
                    # 检查变量是否逃逸（使用 DataFlowInfo 的方法）
                    if dataflow.variable_escapes(var_name):
                        # 变量逃逸：使用堆模式
                        storage_mode = "heap"
                    else:
                        # 变量不逃逸：可以使用栈模式（支持默认参数）
                        storage_mode = "stack"

            # 设置存储模式
            node.set_extension_param("storage_mode", storage_mode)

    def _analyze_async_dataflow(
        self, func: IRFunctionNode, dataflow: DataFlowInfo
    ) -> None:
        """分析异步函数的数据流

        使用新的 AsyncDataFlowAnalyzer 进行分析：
        1. Segment 划分（有父子关系，parent_segment_id）
        2. Unit 划分（无父子关系，记录 segment_id）
        3. 基于 Segment 父链的变量提升分析

        Args:
            func: 异步函数节点
            dataflow: 通用数据流信息
        """
        from mcpy.analysis.async_dataflow import AsyncDataFlowAnalyzer

        # 使用新的分析器
        analyzer = AsyncDataFlowAnalyzer()
        async_info = analyzer.analyze(func)

        # 设置 extension 参数
        func.set_extension_param("async_dataflow", async_info)
